package com.starhub.application.collection.client;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
import com.starhub.application.collection.common.MetaDataLinkHelper;
import com.starhub.application.collection.dto.MetaDataField;
import com.starhub.application.collection.dto.MetaDataLinkDto;
import com.starhub.application.collection.dto.MetaDataSimpleData;
import com.starhub.application.collection.mapper.MetadataRelationMapper;
import com.starhub.application.collection.service.KnowledgeBaseService;
import com.starhub.application.rag.constant.RagConstant;
import com.starhub.application.rag.dto.SegmentDto;
import com.starhub.application.rag.service.EmbeddingService;
import com.starhub.application.rag.service.SegmentService;
import com.starhub.application.rag.store.MilvusStore;
import com.starhub.engine.generator.genid.IdGenerator;
import com.starhub.integration.milvus.langchain.MilvusMetadataFilterMapper;
import com.starhub.utils.Util;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.store.embedding.filter.Filter;
import dev.langchain4j.store.embedding.filter.MetadataFilterBuilder;
import lombok.extern.slf4j.Slf4j;

import com.starhub.application.collection.entity.KnowledgeBase;
import com.starhub.application.collection.entity.MetadataRelation;

/**
 * 关联数据客户端
 * 
 * @version: V1.0
 * @author: liuhf
 * @time 2025-06-09
 */
@Slf4j
@Component
public class LinkDataClient {

    @Autowired
    private MetadataRelationMapper metadataRelationMapper;

    @Autowired
    private MilvusStore milvusStore;

    @Autowired
    private MetaDataLinkHelper metaDataLinkHelper;

    @Autowired
    private SegmentService segmentService;

    @Autowired
    private KnowledgeBaseService knowledgeBaseService;

    @Autowired
    private EmbeddingService embeddingService;

    /**
     * 添加关联数据
     * 
     * @param metaDataLinkDto
     * @throws Exception
     */
    public boolean addTargetLinkData(MetaDataLinkDto metaDataLinkDto) throws Exception {
        // 创建MetaDataSimpleData列表
        /*
         * List<MetaDataSimpleData> metaDataList =
         * metaDataLinkDto.getSelectIds().stream()
         * .map(id -> MetaDataSimpleData.builder()
         * .id(id)
         * .collectionId(metaDataLinkDto.getMetaCollectionId())
         * .build())
         * .collect(Collectors.toList());
         * 
         * metaDataLinkDto.setMetaDatas(metaDataList);
         */

        // 1. 基础数据
        String itemId = metaDataLinkDto.getItemId();
        List<MetaDataSimpleData> metaDatas = metaDataLinkDto.getMetaDatas();
        String collectionName = metaDataLinkDto.getCollectionIdentifier();

        // 2. 获取现有关联数据
        LambdaQueryWrapper<MetadataRelation> queryWrapper = new LambdaQueryWrapper<>();
        queryWrapper.eq(MetadataRelation::getMainDataId, Util.getLongValue(itemId))
                .eq(MetadataRelation::getMetaCollectionId, metaDataLinkDto.getMetaCollectionId());
        List<MetadataRelation> relations = metadataRelationMapper.selectList(queryWrapper);

        // 3. 过滤出需要新增的关联数据
        List<MetadataRelation> newRelations = metaDatas.stream()
                .filter(meta -> relations.stream()
                        .noneMatch(relation -> relation.getMetaDataId().equals(meta.getId())))
                .map(meta -> {
                    MetadataRelation relation = new MetadataRelation();
                    relation.setId(IdGenerator.generateId());
                    relation.setMainDataId(itemId);
                    relation.setMainCollectionId(Util.getLongValue(metaDataLinkDto.getCollectionId()));
                    relation.setMainBaseId(Util.getLongValue(metaDataLinkDto.getBaseId()));
                    relation.setMetaCollectionId(Util.getLongValue(metaDataLinkDto.getMetaCollectionId()));
                    relation.setMetaDataId(meta.getId());
                    metadataRelationMapper.insert(relation);
                    return relation;
                })
                .collect(Collectors.toList());

        // 4. 合并现有关联和新增关联
        List<MetadataRelation> allRelations = Stream.concat(relations.stream(), newRelations.stream())
                .collect(Collectors.toList());

        // 5. 更新Milvus数据
        MetaDataField metaDataField = metaDataLinkHelper.getMetaDataLinkField(allRelations,
                metaDataLinkDto.getMetaCollectionIdentifier());

        //updateMetadata(metaDataField, collectionName, itemId);
        updateLinkData(metaDataField, metaDataLinkDto);
        return true;
    }

    /**
     * 更新关联数据
     * 
     * @param metaDataLinkDto
     * @throws Exception
     */
    public boolean updateTargetLinkData(MetaDataLinkDto metaDataLinkDto) throws Exception {
        // 1. 基础数据
        String itemId = metaDataLinkDto.getItemId();
        List<MetaDataSimpleData> metaDatas = metaDataLinkDto.getMetaDatas();
        String collectionName = metaDataLinkDto.getCollectionIdentifier();

        // 2. 删除现有关联数据
        LambdaQueryWrapper<MetadataRelation> queryWrapper = new LambdaQueryWrapper<>();
        queryWrapper.eq(MetadataRelation::getMainDataId, Util.getLongValue(itemId))
                .eq(MetadataRelation::getMetaCollectionId, metaDataLinkDto.getMetaCollectionId());
        metadataRelationMapper.delete(queryWrapper);

        // 3. 过滤出需要新增的关联数据
        List<MetadataRelation> relations = metaDatas.stream()
                .map(meta -> {
                    MetadataRelation relation = new MetadataRelation();
                    relation.setId(IdGenerator.generateId());
                    relation.setMainDataId(itemId);
                    relation.setMainCollectionId(Util.getLongValue(metaDataLinkDto.getCollectionId()));
                    relation.setMainBaseId(Util.getLongValue(metaDataLinkDto.getBaseId()));
                    relation.setMetaCollectionId(Util.getLongValue(metaDataLinkDto.getMetaCollectionId()));
                    relation.setMetaDataId(meta.getId());
                    metadataRelationMapper.insert(relation);
                    return relation;
                })
                .collect(Collectors.toList());

        // 4. 更新Milvus数据
        MetaDataField metaDataField = metaDataLinkHelper.getMetaDataLinkField(relations,
                metaDataLinkDto.getMetaCollectionIdentifier());

        //updateMetadata(metaDataField, collectionName, itemId);
        updateLinkData(metaDataField, metaDataLinkDto);
        return true;
    }

    /**
     * 移除关联数据
     * 
     * @param metaDataLinkDto
     * @throws Exception
     */
    public boolean removeTargetLinkData(MetaDataLinkDto metaDataLinkDto) throws Exception {
        // 1. 基础数据
        String itemId = metaDataLinkDto.getItemId();
        List<MetaDataSimpleData> metaDatas = metaDataLinkDto.getMetaDatas();
        String collectionName = metaDataLinkDto.getCollectionIdentifier();

        // 2. 删除现有关联中跟metaDataLinkDto有关的
        LambdaQueryWrapper<MetadataRelation> queryWrapper = new LambdaQueryWrapper<>();
        queryWrapper.eq(MetadataRelation::getMainDataId, Util.getLongValue(itemId))
                .eq(MetadataRelation::getMetaCollectionId, metaDataLinkDto.getMetaCollectionId());
        if (!CollectionUtils.isEmpty(metaDatas)) {
            List<String> metaDataIds = metaDatas.stream()
                    .map(MetaDataSimpleData::getId)
                    .collect(Collectors.toList());
            queryWrapper.in(MetadataRelation::getMetaDataId, metaDataIds);
        }
        metadataRelationMapper.delete(queryWrapper);

        // 3. 获取现有关联数据
        queryWrapper = new LambdaQueryWrapper<>();
        queryWrapper.eq(MetadataRelation::getMainDataId, Util.getLongValue(itemId))
                .eq(MetadataRelation::getMetaCollectionId, metaDataLinkDto.getMetaCollectionId());
        List<MetadataRelation> relations = metadataRelationMapper.selectList(queryWrapper);

        // 4. 更新Milvus数据
        MetaDataField metaDataField = metaDataLinkHelper.getMetaDataLinkField(relations,
                metaDataLinkDto.getMetaCollectionIdentifier());

        //updateMetadata(metaDataField, collectionName, itemId);
        updateLinkData(metaDataField, metaDataLinkDto);
        return true;
    }

    /**
     * 移除所有关联数据
     * 
     * @param metaDataLinkDto
          * @throws Exception 
          */
         public boolean removeAllTargetLinkData(MetaDataLinkDto metaDataLinkDto) throws Exception {
        // 1. 基础数据
        String itemId = metaDataLinkDto.getItemId();
        String collectionName = metaDataLinkDto.getCollectionIdentifier();

        // 2. 删除现有关联中跟metaDataLinkDto有关的
        LambdaQueryWrapper<MetadataRelation> queryWrapper = new LambdaQueryWrapper<>();
        queryWrapper.eq(MetadataRelation::getMainDataId, Util.getLongValue(itemId))
                .eq(MetadataRelation::getMetaCollectionId, metaDataLinkDto.getMetaCollectionId());
        metadataRelationMapper.delete(queryWrapper);

        // 4. 更新Milvus数据
        MetaDataField metaDataField = metaDataLinkHelper.getMetaDataLinkField(null,
                metaDataLinkDto.getMetaCollectionIdentifier());
        //updateMetadata(metaDataField, collectionName, itemId);
        updateLinkData(metaDataField, metaDataLinkDto);
        return true;
    }

    /**
     * 更新Milvus数据（主数据集合对应的metadata）
     * 
     * @param metaDataField  元数据关联字段
     * @param collectionName 知识库表名
     * @param itemId         主数据ID
     * @return 是否更新成功
     */
    private boolean updateMetadata(MetaDataField metaDataField, String collectionName, String itemId) {
        boolean success = false;
        try {
            // 1. 元数据关联字段
            if (metaDataField == null) {
                return true;
            }

            // 2. 构建查询条件 - 查询metadata中item_id字段等于指定值的记录
            Filter filter = new MetadataFilterBuilder(RagConstant.EmbeddingMetaData.ITEM_ID).isEqualTo(itemId);
            String condition = MilvusMetadataFilterMapper.map(filter, RagConstant.LangchainMilvusFields.FIELD_METADATA);

            // 3. 查询所有匹配的数据（主数据的多个片段）
            List<MilvusStore.DocumentWithMetadata> documents = milvusStore.queryByCondition(condition, collectionName);

            if (!documents.isEmpty()) {
                // 4. 批量更新metadata字段
                Map<String, Metadata> updateMap = new HashMap<>();

                List<SegmentDto> segmentDtos = new ArrayList<>();

                for (MilvusStore.DocumentWithMetadata doc : documents) {
                    // 获取现有metadata
                    Metadata existingMetadata = doc.getTextSegment().metadata();

                    // 创建新的metadata，保留原有字段并更新指定字段
                    Map<String, Object> metadataMap = new HashMap<>(existingMetadata.toMap());
                    metadataMap.put(metaDataField.getFieldName(), metaDataField.getFieldValue());

                    Metadata updatedMetadata = Metadata.from(metadataMap);
                    updateMap.put(doc.getId(), updatedMetadata);

                    segmentDtos.add(
                            SegmentDto.builder()
                                    .id(Util.null2String(metadataMap.get(RagConstant.EmbeddingMetaData.SEGMENT_ID)))
                                    .segmentMetadata(updatedMetadata.toString())
                                    .build());
                }
                // 更新片段对应数据
                segmentService.batchUpdate(segmentDtos);
                // 4. 执行批量更新
                success = milvusStore.updateMetadataBatch(updateMap, collectionName);
            }

        } catch (Exception e) {
            log.error("更新Milvus数据失败", e);
            throw new RuntimeException("更新Milvus数据失败: " + e.getMessage(), e);
        }
        return success;
    }

    /**
     * 更新关联集合数据（更新关联数据集合）
     * 
     * @param relations
     * @param collectionName
     * @param itemId
     * @param metacollectionName
     * @return
     * @throws Exception
     */
    private void updateLinkData(MetaDataField metaDataField, MetaDataLinkDto metaDataLinkDto) throws Exception {
        // 1. 获取知识库
        KnowledgeBase knowledgeBase = knowledgeBaseService
                .selectByIdAndNotDeleted(Util.getLongValue(metaDataLinkDto.getBaseId()));
        if (knowledgeBase == null) {
            throw new IllegalArgumentException("所属知识库不存在");
        }

        String collectionName = metaDataLinkDto.getCollectionIdentifier() + RagConstant.META_POSTFIX;

        // 2 删除现有关联数据
        Filter filter1 = new MetadataFilterBuilder(RagConstant.LinkMetaData.MAIN_ID)
                .isEqualTo(metaDataLinkDto.getItemId());
        String condition1 = MilvusMetadataFilterMapper.map(filter1, RagConstant.LangchainMilvusFields.FIELD_METADATA);

        Filter filter2 = new MetadataFilterBuilder(RagConstant.LinkMetaData.META_COLLECTION)
                .isEqualTo(metaDataLinkDto.getMetaCollectionId());
        String condition2 = MilvusMetadataFilterMapper.map(filter2, RagConstant.LangchainMilvusFields.FIELD_METADATA);

        String condition = condition1 + " && " + condition2;

        milvusStore.deleteByCondition(condition, collectionName);

        if(Util.isEmpty(metaDataField.getFieldValue())){
            log.warn("关联数据为空，不进行向量存储 {},{}", metaDataLinkDto.getItemId(), metaDataLinkDto.getMetaCollectionId());
            return;
        }
        
        // 3. 构建metadata
        Metadata metadata = new Metadata();
        metadata.put(RagConstant.LinkMetaData.MAIN_ID, Util.null2String(metaDataLinkDto.getItemId()));
        metadata.put(RagConstant.LinkMetaData.MAIN_COLLECTION, Util.null2String(metaDataLinkDto.getCollectionId()));
        metadata.put(RagConstant.LinkMetaData.META_COLLECTION,
                Util.null2String(metaDataLinkDto.getMetaCollectionId()));
        metadata.put(RagConstant.LinkMetaData.BASE_ID, Util.null2String(metaDataLinkDto.getBaseId()));

        // 4. 关联数据向量化
        Embedding embedding = embeddingService.embedding(knowledgeBase.getVectorModel(),
                Util.null2String(metaDataField.getFieldValue()));

        // 5. 向量存储
        milvusStore.add(Util.null2String(metaDataField.getFieldValue()), embedding, collectionName, metadata);
    }
}
